Puzzle 12 - compress解法

def compress_spec(g, v, out):
    j = 0
    for i in range(len(g)):
        if g[i]:
            out[j] = v[i]
            j += 1

def compress(g: TT["i", bool], v: TT["i"], i:int) -> TT["i"]:
    
    raise NotImplementedError


test_compress = make_test("compress", compress, compress_spec, add_sizes=["i"])

核心目标与思路

给定布尔掩码 $g$ 和数据张量 $v$ ：

1.目标是将所有满足 $g[i] == \text{True}$ 的元素 $v[i]$ 依次排在输出的前面。

2.在纯张量无循环规则下，我们需要找到每个有效元素在新张量中的目标下标 $j$ 。

关键观察：

$v[i]$ 如果被选中，它在新张量里的索引 $j$ 就是它前面（包含自身）为True的元素总个数减 1，即 前缀和（cumsum）。
记 $c = \text{cumsum}(g) - 1$ 。当且仅当 $g[i] == \text{True}$ 时， $v[i]$ 会被放置在索引 $c[i]$ 处。
用外积/广播构造比较矩阵：看目标索引矩阵（形状 $[i]$ ）与各个位置计算出的目标索引何时匹配

1.构造前缀和索引：对于长度为 $i$ 的序列，构造二维矩阵比较：

# j <= i 形成的下三角布尔矩阵
(arange(i)[:, None] <= arange(i)[None, :])

将元素散射（Scatter）到对应输出位置：

设目标位置索引为 $k \in [0, i)$ ：

位置 $k$ 的值来自于满足g[m] == True且 prefix_sum[m] == k 的那个原位置 $m$ 。

匹配条件：(prefix_sum[:, None] == arange(i)[None, :]) & g[:, None]
利用这个掩码对原输入 $v$ 进行加权求和即可将元素放置到指定下标。

方案一：使用内置Cumsum

def compress(g: TT["i", bool], v: TT["i"], i: int) -> TT["i"]:
    # 1. 目标下标为前缀和 - 1
    target_idx = g.int().cumsum(0) - 1
    
    # 2. 判断原数组中的位置与输出数组中的位置是否匹配且处于激活状态
    # mask 形状: (i, i)，第 row 行对应原索引，第 col 列对应输出位置
    match = (target_idx[:, None] == arange(i)[None, :]) & g[:, None]
    
    # 3. 聚合并输出
    return (v[:, None] * match).sum(0)

评论 0